import numpy as np
import matplotlib.pyplot as plt

np.random.seed(1)

z = [i for i in range(100)]
z_watch = np.mat(z)

noise = np.round(np.random.normal(0, 1, 100), 2)
noise_mat = np.mat(noise)

z_mat = z_watch + noise_mat

print('z_mat', z_mat.shape)
print(z_mat)

d_s = np.concatenate([z_mat[:, 0:1], z_mat[:, 1:] - z_mat[:, :-1]], axis=1)
d_t = 1.
v_mat = d_s / d_t
print('v_mat', v_mat.shape)
print('mean v:', v_mat.mean())

# 初始状态 s=0 v=0
x_mat = np.mat([[0], [0]])  # 2x1
print('x_mat:\n', x_mat)
# 初始状态协方差矩阵
p_mat = np.mat([[1, 0], [0, 1]])  # 2x2
print('p_mat:\n', p_mat)

# 状态转移矩阵
a_mat = np.mat([[1, d_t], [0, 1]])   # 2x2

# 定义状态转移协方差矩阵，这里我们把协方差设置的很小，因为觉得状态转移矩阵准确度高
PROCESS_NOISE = 1e-4
q_mat = np.mat([[PROCESS_NOISE, 0], [0, PROCESS_NOISE]])

# 观测矩阵 观测s 观测不到v
h_mat = np.mat([1, 0])  # 1x2

# 定义观测噪声协方差
r_mat = np.mat([1])

for i in range(z_mat.size):
    x_pred = a_mat * x_mat  # 2x1
    p_pred = a_mat * p_mat * a_mat.T + q_mat  # 2x2 * 2x2 * 2x2 => 2x2 + 2x2 => 2x2

    # (2x2 * 2x1 => 2x1) / (1x2 * 2x2 * 2x1 => 1x1 + 1x1 => 1x1) => 2x1
    kalman = (p_pred * h_mat.T) / (h_mat * p_pred * h_mat.T + r_mat)

    # 2x1 + [2x1 * (1x1 - 1x2 * 2x1 => 1x1 => 1x1) => 2x1] => 2x1
    x_mat = x_pred + kalman * (z_mat[0, i] - h_mat * x_pred)
    plt.plot(x_mat[0, 0], x_mat[1, 0], 'ro', markersize=2)

    # [(2x2 - 2x1 * 1x2 => 2x2) => 2x2 ] * 2x2 => 2x2
    p_mat = (np.eye(2) - kalman * h_mat) * p_pred

    if i == 0:
        print('x_pred', x_pred.shape)
        print('p_pred', p_pred.shape)
        print('kalman', kalman.shape)
        print('x_mat', x_mat.shape)

plt.plot(np.array(v_mat).ravel(), label='velocity')
plt.legend()
plt.show()
